#!/usr/bin/env python
import threading
import struct
import socket
import select
import string
from observable import Observable
from Queue import Queue

SOCKS_VERSION = 4
SOCKS_CONNECT = 1
SOCKS_BIND = 2
SOCKS_REQUEST_GRANTED                 = 90
SOCKS_REQUEST_REJECTED_FAILED         = 91
SOCKS_REQUEST_REJECTED_NO_IDENTD      = 92
SOCKS_REQUEST_REJECTED_IDENT_FAILED   = 93

FILTER=''.join([(len(repr(chr(x)))==3) and chr(x) or '.' for x in range(256)])

def hex(src, length=12):
    result=[]
    for i in xrange(0, len(src), length):
        s = src[i:i+length]
        hexa = ' '.join(["%02X"%ord(x) for x in s])
        printable = s.translate(FILTER)
        #result.append("%04X   %-*s   %s\n" % (i, length*3, hexa, printable))
        result.append("%04X   %-*s   %s\n" % (i, length*3, hexa, printable))
    return ''.join(result)
    

class Socks4Server(Observable):
    def __init__(self, config):
        Observable.__init__(self)
        self.config = config
        self.running = True
        self.messages = Queue()
        self.edit = False
        self.intercept_in = False
        self.intercept_out = False

    def listen(self):
        try:
            self.ls = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            self.ls.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
            self.ls.bind((self.config["socks_bind_address"], self.config["socks_port"]))
            self.ls.listen(5)
        except Exception, e:
            self.notifyg({"event":"exception", "exception":e})
            return
            
        while self.running:
            try:
                self.cs,a = self.ls.accept()
                data = self.cs.recv(self.config["socks_buffer_size"])
                r = self.parse_request(data)
                if r["cd"] == SOCKS_CONNECT:
                    self.handle_connect(r)
                elif r["cd"] == SOCKS_BIND:
                    self.handle_bind(r)
                self.notifyg({"event":"client_connected", "client":self})
                self.handle() 
            except Exception, e:
                self.notifyg({"event":"exception", "exception":e})
        
        self.ls.close()
        self.cs.close()
        self.ss.close()

            
    def parse_request(self, data):
        if len(data) < 9 or "\x00" not in data: 
            raise Exception("invalid message.", data)
        r = {}
        r["version"],r["cd"],r["port"]=struct.unpack("!BBH",data[:4])
        if r["version"] != 4 or r["cd"] not in (1,2) or r["port"] not in range(0, 65536):
            raise Exception("invalid data.", data, r["version"], r["cd"], r["port"])
        if data[4:7] == "\x00\x00\x00":
            dns = string.split(data[9:], '\x00', 1)[0]
            try:
                r["server"] = socket.gethostbyname(dns)
            except socket.error:
                raise Exception("invalid address.", data, r["server"])
        else:
            try:
                r["server"] = socket.inet_ntoa(data[4:8])
            except socket.error:
                raise Exception("invalid address.", data, r["server"])
            
        return r
    
    def handle_connect(self, r):
        self.ss = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        self.ss.settimeout(self.config["socks_server_timeout"])
        try:               
            self.ss.connect((r["server"], r["port"]))
        except socket.error:
            self.answer(SOCKS_REQUEST_REJECTED_FAILED, r)
            raise Exception("connection error.")
        try:
            self.answer(SOCKS_REQUEST_GRANTED, r)
        except socket.error:
            raise Exception("socket error.")

    def answer(self, cd, r):
        try:
            p = struct.pack('!BBH', 0, cd, r["port"])
            p += socket.inet_aton(r["server"])
            self.cs.send(p)
        except:
            raise Exception("client closed connection")
                    
    def handle(self):
        try:
            socketlist = [self.ss, self.cs]
            while True:
                r, w, x = select.select(socketlist, [], [], self.config["socks_client_timeout"])
                if x or (not r and not w and not x):
                    raise Exception("client timeout")
                buffer = ""
                for s in r:
                    ws = [self.ss, self.cs]
                    ws.remove(s)
                    o = ws[0]
                    buffer = s.recv(self.config["socks_buffer_size"])
                    if buffer:
                        if not self.intercept_in and not self.intercept_out:
                            o.send(buffer)
                        elif self.intercept_in and o is self.cs:
                            self.messages.put([o, buffer])
                            self.notify({"event":"msg_received"})
                        elif self.intercept_out and o is self.ss:
                            self.messages.put([o, buffer])
                            self.notify({"event":"msg_received"})
                    else:
                        raise Exception("client closed connection")
        except socket.error:
            raise 
        
    def send(self, socket, message):
        self.messages.task_done()
        try:
            socket.send(message)
        except Exception, e:
            self.notifyg({"event":"exception", "exception":e})
        if not self.messages.empty():
            self.notifyg({"event":"msg_received"})

    def start(self):
        t = threading.Thread(target = self.listen)
        t.setDaemon(True)
        t.start()
            
    def stop(self):
        self.running = False       
               
    def pass_event(self, eventdict):
        self.notifyg(eventdict)
                    
    def drop(self, message):
        self.client.messages.task_done()   
